#include "UDParser.h"

#include <clang/Frontend/TextDiagnosticBuffer.h>

#ifdef DEBUG
#include <iostream>
#endif
#include <string>

#include "IncrementalParser.h"
#include "common/Coder.h"
#include "common/Searcher.h"

using namespace rocket::incparser;
using namespace clang;
using namespace std;

static const string TmpSpecToolFuncName = "__rocket_TEMPLATE_TYPE_IC";
static const string TmpSpecToolRecordName = "__rocket_TEMPLATE_RECORD_IC";

static string getDeclTmpName() {
  static unsigned int count = 0;
  return "__rocket_INCREMENTAL_PARSER_TEMP_VARABLE_NAME_" + to_string(count++);
};

UDParser::UDParser(CompilerInstance *_CI) : CI(_CI) {
  CI->getPreprocessor().enableIncrementalProcessing();
  if (CI->getLangOpts().CPlusPlus) {
    IncrementalParser parser(*CI);
    TextDiagnosticBuffer diag_buf;
    string func =
        "template<typename T, int count> void " + TmpSpecToolFuncName + "(){}";

    // TODO: 模板形式的声明方案
    string record = "template<typename T> struct " + TmpSpecToolRecordName +
                    " { typedef T TYPE; };";

    string errors;
    parser.parse(func, diag_buf)
        .onTopLevelDecl([this](DeclGroupRef D) {
          for (DeclGroupRef::iterator it = D.begin(), e = D.end(); it != e;
               ++it) {
            if (FunctionTemplateDecl *ftd =
                    dyn_cast<FunctionTemplateDecl>(*it)) {
              if (rocket::common::getFullNDeclName(ftd) ==
                  TmpSpecToolFuncName) {
                this->FTD = ftd;
                return false;
              }
            }
          }
          return true;
        })
        .onError([this, &diag_buf, &errors]() {
          const SourceManager &sm = CI->getSourceManager();
          for (auto it = diag_buf.err_begin(); it != diag_buf.err_end(); ++it) {
            PresumedLoc presumed_loc = sm.getPresumedLoc(it->first);
            if (presumed_loc.isValid())
              errors += it->second + ", at " + ":" +
                        to_string(presumed_loc.getLine()) + ":" +
                        to_string(presumed_loc.getColumn()) + "\n";
          }
        });
    if (!errors.empty()) {
      throw runtime_error("CompilerInstance build error:" + errors);
    }

    //    CI->getSema().getDiagnostics().Reset();
    flushASTContext();
  } else {
    FTD = nullptr;
  }
}

const string rocket::incparser::UDParser::flushASTContext() const {
  // FIXME:函数模板特化产生Diagnotics时会导致Act内产生运行时错误
  CI->getSema().getDiagnostics().setIgnoreAllWarnings(true);
  CI->getSema().ActOnEndOfTranslationUnit();
  CI->getSema().getDiagnostics().Reset();
  CI->getSema().getDiagnostics().setIgnoreAllWarnings(false);
  return "";
}

QualType rocket::incparser::UDParser::parseType(const string &var_name,
                                                const string &var_decl) {
  // 准备增量编译
  IncrementalParser parser(*CI);
  TextDiagnosticBuffer diag_buf;
  QualType type;
  string errors;
  parser.parse(var_name, var_decl, diag_buf)
      .onTopLevelDecl([&type, &var_name](DeclGroupRef D) {
        for (DeclGroupRef::iterator it = D.begin(), e = D.end(); it != e;
             ++it) {
          if (DeclaratorDecl *vd = dyn_cast<DeclaratorDecl>(*it)) {
            if (vd->getName() == var_name) {
              type = vd->getType();
              return false;
            }
          }
        }
        return true;
      })
      .onError([this, &diag_buf, &var_decl, &errors]() {
        const SourceManager &sm = CI->getSourceManager();
        for (auto it = diag_buf.err_begin(); it != diag_buf.err_end(); ++it) {
          PresumedLoc presumed_loc = sm.getPresumedLoc(it->first);
          if (presumed_loc.isValid())
            errors += it->second + ", at " + var_decl + ":" +
                      to_string(presumed_loc.getLine()) + ":" +
                      to_string(presumed_loc.getColumn()) + "\n";
        }
      });
  if (type.isNull()) {
    throw IllegalTypeParseException("failed to parse type : " + errors);
  } else {
    if (!errors.empty())
      exceptions.push_back(IllegalTypeParseException(
          "error occured when parsing type : " + errors));
    return type;
  }
}

QualType UDParser::parseType(const string &type_str) {
  if (!FTD) {
    // 声明法获得类型特化
    string name = getDeclTmpName();
    string decl = "extern " + rocket::common::makeVarDecl(type_str, name) + ";";
    return parseType(std::move(name), std::move(decl));
  } else {
    // 模板特化法获取类型特化
    static unsigned int tspec_count = 0;
    std::string tspec_identify = to_string(tspec_count);

    string spec = "template<> void " + TmpSpecToolFuncName + "<" + type_str +
                  "," + tspec_identify + ">(){};";

    IncrementalParser parser(*CI);
    TextDiagnosticBuffer diag_buf;

    string errors;
    parser.parse(spec, diag_buf).onError([this, &diag_buf, &errors]() {
      const SourceManager &sm = CI->getSourceManager();
      for (auto it = diag_buf.err_begin(); it != diag_buf.err_end(); ++it) {
        PresumedLoc presumed_loc = sm.getPresumedLoc(it->first);
        if (presumed_loc.isValid())
          errors += it->second + ", at " + ":" +
                    to_string(presumed_loc.getLine()) + ":" +
                    to_string(presumed_loc.getColumn()) + "\n";
      }
    });

    flushASTContext();

    for (auto spec : FTD->specializations()) {
      const auto &spec_args = spec->getTemplateSpecializationArgs();
      assert(spec_args->size() == 2);
      auto spec_type = spec_args->get(0);
      auto spec_count = spec_args->get(1);

      auto spec_count_str =
          rocket::common::getFullSpecArgStr(spec_count, CI->getASTContext());
      if (spec_count_str == tspec_identify) {
        ++tspec_count;
        if (!errors.empty())
          exceptions.push_back(IllegalTypeParseException(
              "errors occured when parsing type " + type_str + " : " + errors));
        return spec_type.getAsType();
      }
    }
    throw IllegalTypeParseException("failed to parse type " + type_str + " : " +
                                    errors);
  }
}

QualType UDParser::parseTemplateDependentType(
    string type_str, const std::vector<TemplateArgument> &ud_tas) {
  if (!ud_tas.empty()) {
    type_str.append("<");
    for (auto ta : ud_tas) {
      type_str.append(common::getFullSpecArgStr(ta, CI->getASTContext()) + ",");
    }
    type_str.pop_back();
    type_str.append(">");
  }
  return parseType(type_str);
}

vector<string> UDParser::parseTemplateDependentTypeListStrings(
    const TemplateDecl *temp_decl, const std::vector<QualType> type_list,
    const std::vector<TemplateArgument> &ud_tas) {
  vector<string> spec_type_names;

  string temp_rd_name = getUniqueName();

  string temp_rd_scope = temp_rd_name;
  if (!ud_tas.empty()) {
    temp_rd_scope.append("<");
    for (auto &ta : ud_tas) {
      temp_rd_scope.append(common::getFullSpecArgStr(ta, CI->getASTContext()) +
                           ",");
    }
    temp_rd_scope.pop_back();
    temp_rd_scope.append(">");
    temp_rd_scope += "::";
  }

  string temp_rd_code =
      common::makeTemplateParmList(temp_decl->getTemplateParameters()) +
      " struct " + temp_rd_name + "{";

  for (auto &dep_type : type_list) {
    string tmp_type_str = getUniqueName();
    temp_rd_code += "typedef " +
                    common::getFullQTypeName(dep_type, CI->getASTContext()) +
                    " " + tmp_type_str + ";";
    spec_type_names.push_back(temp_rd_scope + tmp_type_str);
  }

  temp_rd_code += "};";
  parseCode(temp_rd_code);

  flushASTContext();
  return spec_type_names;
}

vector<QualType> UDParser::parseTemplateDependentTypeListTypes(
    const TemplateDecl *temp_decl, const std::vector<QualType> type_list,
    const std::vector<TemplateArgument> &ud_tas) {
  vector<QualType> types;
  auto type_strs =
      parseTemplateDependentTypeListStrings(temp_decl, type_list, ud_tas);
  for (auto &type_str : type_strs) {
    types.push_back(parseType(type_str));
  }
  return types;
}

QualType UDParser::parseBasicTypeToReference(QualType type) {
  string name = "&" + getDeclTmpName();
  string decl = "extern " +
                rocket::common::makeVarDecl(type, name, CI->getASTContext()) +
                ";";
  QualType qt = parseType(std::move(name), std::move(decl));
  assert(qt->isReferenceType());
  return qt;
}

QualType UDParser::parseBasicTypeToPointer(QualType type) {
  string name = "*" + getDeclTmpName();
  string decl = "extern " +
                rocket::common::makeVarDecl(type, name, CI->getASTContext()) +
                ";";
  QualType qt = parseType(std::move(name), std::move(decl));
  assert(qt->isPointerType());
  return qt;
}

QualType UDParser::parseBasicTypeToArray(QualType type, unsigned length) {
  string name = getDeclTmpName();
  string decl =
      "extern " +
      rocket::common::makeVarDecl(type, name + "[" + to_string(length) + "]",
                                  CI->getASTContext()) +
      ";";
  QualType qt = parseType(std::move(name), std::move(decl));
  return qt;
}

QualType UDParser::parseFunctionType(QualType ret_type,
                                     const std::vector<QualType> &parm_types) {
  string name = getDeclTmpName();
  string decl = "(*" + name + ")";

  if (parm_types.empty()) {
    decl.append("()");
  } else {
    decl.append("(");
    for (auto &parm_type : parm_types) {
      decl.append(common::getFullQTypeName(parm_type, CI->getASTContext()) +
                  ",");
    }
    decl.pop_back();
    decl.append(")");
  }

  decl = common::makeVarDecl(ret_type, decl, CI->getASTContext());
  QualType qt = parseType(std::move(name), decl + ";");
  if (qt->isPointerType()) {
    qt = qt->getPointeeType();
  }
  return qt;
}

DeclaratorDecl *UDParser::parseInitlization(string var_name, string var_decl) {
  incparser::IncrementalParser parser(*CI);
  clang::TextDiagnosticBuffer diag_buf;
  clang::DeclaratorDecl *decl;
  std::string errors;
  parser.parse(var_name, var_decl, diag_buf)
      .onTopLevelDecl([&decl, &var_name](clang::DeclGroupRef D) {
        for (clang::DeclGroupRef::iterator it = D.begin(), e = D.end(); it != e;
             ++it) {
          if (clang::DeclaratorDecl *vd =
                  llvm::dyn_cast<clang::DeclaratorDecl>(*it)) {
            if (vd->getName() == var_name) {
              decl = vd;
              return false;
            }
          }
        }
        return true;
      })
      .onError([this, &diag_buf, &var_decl, &errors]() {
        const clang::SourceManager &sm = CI->getSourceManager();
        for (auto it = diag_buf.err_begin(); it != diag_buf.err_end(); ++it) {
          clang::PresumedLoc presumed_loc = sm.getPresumedLoc(it->first);
          if (presumed_loc.isValid())
            errors += it->second + ", at " + var_decl + ":" +
                      std::to_string(presumed_loc.getLine()) + ":" +
                      std::to_string(presumed_loc.getColumn()) + "\n";
        }
      });
  if (!errors.empty()) {
    throw IllegalTypeParseException("failed to parse expr : " + errors);
  }
  return decl;
}

TemplateArgument UDParser::parseTemplateArgument(
    const shared_ptr<UDParser::UDTemplateArg> &ta) {
  if (auto type_arg = dynamic_pointer_cast<UDTTypeArg>(ta))
    return parseTypeTemplateArgument(type_arg->type);
  else if (auto nontype_arg = dynamic_pointer_cast<UDTNonTypeArg>(ta))
    return parseNonTypeTemplateArgument(nontype_arg->type, nontype_arg->expr);
  else if (auto tmptmp_arg = dynamic_pointer_cast<UDTTempArg>(ta))
    return parseTemplateTemplateArgument(tmptmp_arg->name);
  else
    throw /* impossible */;
}

TemplateArgument UDParser::parseTypeTemplateArgument(string type) {
  QualType tta = parseType(type);
  return TemplateArgument(tta);
}

TemplateArgument UDParser::parseNonTypeTemplateArgument(string type,
                                                        string expr) {
  QualType ntta_t = parseType(type);

  // TODO:模板实参的形式构建表达式
  std::string ntta_name =
      "__rocket_NON_TYPE_TEMPLATE_ARGUEMENT_NAME_" + getUniqueName();
  auto init_decl = dyn_cast<VarDecl>(parseInitlization(
      ntta_name, common::makeVarDecl(ntta_t, ntta_name, CI->getASTContext()) +
                     "=" + expr + ";"));
  assert(init_decl);
  auto *init_expr = init_decl->getInit();
  assert(init_expr);

  if (IntegerLiteral *il = dyn_cast<IntegerLiteral>(init_expr)) {
    // 整数
    return TemplateArgument(il);
  } else {
    // 其他表达式
    return TemplateArgument(init_expr);
  }
}

TemplateArgument UDParser::parseTemplateTemplateArgument(string temp) {
  // TODO:可能的获取ClassTemplateDecl的方法
  clang::ClassTemplateDecl *ctd =
      common::findNamedCXXRecord(CI->getASTContext().getTranslationUnitDecl(),
                                 temp)
          ->getDescribedClassTemplate();
  if (!ctd) {
    throw runtime_error("could not found the ClassTemplateDecl named " + temp);
  }
  return TemplateArgument(TemplateName(ctd));
}

QualType UDParser::specCXXRecordTemplate(
    string rd_name, const std::vector<TemplateArgument> &tas, bool mem_spec) {
  auto rd_decl = common::findNamedCXXRecord(
      CI->getASTContext().getTranslationUnitDecl(), rd_name);
  if (!rd_decl) {
    throw IllegalTypeParseException("could not found CXXRecordDecl named " +
                                    rd_name);
  }
  return specCXXRecordTemplate(rd_decl, tas, mem_spec);
}

QualType UDParser::specCXXRecordTemplate(
    string rd_name,
    const std::vector<std::shared_ptr<UDParser::UDTemplateArg>> &tas,
    bool mem_spec) {
  std::vector<TemplateArgument> args;
  for (auto &ta : tas) {
    args.push_back(parseTemplateArgument(ta));
  }

  return specCXXRecordTemplate(rd_name, args, mem_spec);
}

QualType UDParser::specCXXRecordTemplate(
    const CXXRecordDecl *rd, const std::vector<TemplateArgument> &tas,
    bool mem_spec) {
  if (!rd->isTemplated()) return QualType(rd->getTypeForDecl(), 0);

  if (auto rtd = rd->getDescribedClassTemplate()) {
    return specCXXRecordTemplate(rtd, tas, mem_spec);
  } else {
    return QualType();
  }
}

QualType UDParser::specCXXRecordTemplate(
    const ClassTemplateDecl *ctd, const std::vector<TemplateArgument> &tas,
    bool mem_spec) {
  if (!ctd->isTemplated())
    return QualType(ctd->getTemplatedDecl()->getTypeForDecl(), 0);

  string rd_name = common::getFullNDeclName(ctd->getTemplatedDecl());

  QualType rd_sepc_type = parseTemplateDependentType(rd_name, tas);

  flushASTContext();

  if (mem_spec && rd_sepc_type->isRecordType()) {
    auto rd_decl = rd_sepc_type->getAsRecordDecl();
    if (ClassTemplateSpecializationDecl *rd_spec_decl =
            dyn_cast<ClassTemplateSpecializationDecl>(rd_decl)) {
      ClassTemplateDecl *ctd = rd_spec_decl->getSpecializedTemplate();
      auto tas = rd_spec_decl->getTemplateArgs().asArray();
      auto rd_temp_decl = ctd->getTemplatedDecl();
      if (rd_temp_decl) {
        // 特化成员函数
        for (auto fd : rd_temp_decl->methods()) {
          if (isa<CXXConstructorDecl>(fd)) continue;

          if (fd->isTemplateDecl()) {
            // 仍为模板函数，只需声明或原型，不需要将其特化
          } else {
            /* 函数指针名构建 */
            string fptr_name = getUniqueName();

            // 构建成员指针式的临时特化变量
            fptr_name =
                common::getFullQTypeName(rd_sepc_type, CI->getASTContext()) +
                "::*" + fptr_name;

            vector<QualType> type_list;
            type_list.push_back(fd->getReturnType());

            for (auto param : fd->parameters()) {
              type_list.push_back(param->getType());
            }

            try {
              /* 模板函数形参不完整类型特化 与 声明语句构建 */
              std::string incdecl_str = "auto " + fptr_name;

              /* 模板函数模板实参填入 */
              std::string func_name =
                  common::getFullQTypeName(rd_sepc_type, CI->getASTContext()) +
                  "::" + fd->getNameAsString();

              /* 声明语句增量编译实现特化 */
              parseCode(incdecl_str + "=&" + func_name + ";");
            } catch (IllegalTypeParseException &e) {
              exceptions.push_back(e);
            }

            flushASTContext();
          }
        }
      }
    }
  }

  return rd_sepc_type;
}

bool UDParser::specFunctionTemplate(
    const string &fmangled,
    const std::vector<std::shared_ptr<UDParser::UDTemplateArg>> &tas) {
  auto func_decl = common::findMangledFunction(
      CI->getASTContext().getTranslationUnitDecl(), fmangled);
  std::vector<TemplateArgument> args;
  for (auto &ta : tas) {
    args.push_back(parseTemplateArgument(ta));
  }
  return specFunctionTemplate(func_decl, args);
}

bool UDParser::specFunctionTemplate(const string &fmangled,
                                    const std::vector<TemplateArgument> &tas) {
  auto fd = common::findMangledFunction(
      CI->getASTContext().getTranslationUnitDecl(), fmangled);
  if (!fd) {
    throw IllegalTypeParseException(
        "could not found FuncitonDecl with mangled name " + fmangled);
  }
  return specFunctionTemplate(fd, tas);
}

bool UDParser::specFunctionTemplate(const FunctionDecl *fd,
                                    const std::vector<TemplateArgument> &tas) {
  if (!fd) return false;

  if (auto ftd = fd->getPrimaryTemplate()) {
    return specFunctionTemplate(ftd, tas);
  } else if (auto ftd = fd->getDescribedFunctionTemplate()) {
    return specFunctionTemplate(ftd, tas);
  } else {
    return false;
  }
}

bool UDParser::specFunctionTemplate(const FunctionTemplateDecl *ftd,
                                    const std::vector<TemplateArgument> &tas) {
  if (!ftd) return false;

  FunctionDecl *fd = ftd->getTemplatedDecl();
  string fptr_name = getUniqueName();

  /* 函数指针名构建 */
  if (auto md = dyn_cast<CXXMethodDecl>(fd)) {
    auto parent_decl = md->getParent();
    if (parent_decl->isTemplated()) {
      // 所属类未特化时函数不能特化
      return false;
    }

    // 构建成员指针式的临时特化变量
    fptr_name =
        common::getFullQTypeName(QualType(parent_decl->getTypeForDecl(), 0),
                                 CI->getASTContext()) +
        "::*" + fptr_name;
  } else {
    // 构建函数指针式的临时特化变量
    fptr_name = "*" + fptr_name;
  }

  vector<QualType> type_list;
  type_list.push_back(fd->getReturnType());

  for (auto param : fd->parameters()) {
    type_list.push_back(param->getType());
  }

  /* 模板函数形参不完整类型特化 与 声明语句构建 */
  std::string incdecl_str = "auto " + fptr_name;

  /* 模板函数模板实参填入 */
  std::string func_name = common::getFullNDeclName(fd, false);
  if (tas.size() > 0) {
    func_name.append("<");
    for (auto targ : tas) {
      func_name.append(common::getFullSpecArgStr(targ, CI->getASTContext()) +
                       ",");
    }
    func_name.pop_back();
    func_name.append(">");
  }

  /* 声明语句增量编译实现特化 */

  try {
    if (isa<CXXMethodDecl>(fd))
      parseCode(incdecl_str + "=&" + func_name + ";\n");
    else
      parseCode(incdecl_str + "=" + func_name + ";\n");
  } catch (IllegalTypeParseException &e) {
    exceptions.push_back(e);
  }

  flushASTContext();
  return true;
}

string UDParser::getUniqueName() const { return getDeclTmpName(); }

bool UDParser::parseCode(const string &code) {
  IncrementalParser parser(*CI);
  TextDiagnosticBuffer diag_buf;
  string errors;
  parser.parse(code, diag_buf).onError([this, &diag_buf, &errors]() {
    const SourceManager &sm = CI->getSourceManager();
    for (auto it = diag_buf.err_begin(); it != diag_buf.err_end(); ++it) {
      PresumedLoc presumed_loc = sm.getPresumedLoc(it->first);
      if (presumed_loc.isValid())
        errors += it->second + ", at " + ":" +
                  to_string(presumed_loc.getLine()) + ":" +
                  to_string(presumed_loc.getColumn()) + "\n";
    }
  });
  if (!errors.empty()) {
    exceptions.push_back(IllegalTypeParseException(errors));
    return false;
  }
  return true;
}
